{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "import statsmodels.formula.api as smf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>year</th>\n",
       "      <th>state</th>\n",
       "      <th>legal1820</th>\n",
       "      <th>dtype</th>\n",
       "      <th>agegr</th>\n",
       "      <th>count</th>\n",
       "      <th>pop</th>\n",
       "      <th>age</th>\n",
       "      <th>legal</th>\n",
       "      <th>beertaxa</th>\n",
       "      <th>beerpercap</th>\n",
       "      <th>winepercap</th>\n",
       "      <th>spiritpercap</th>\n",
       "      <th>totpercap</th>\n",
       "      <th>mrate</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1970</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>all</td>\n",
       "      <td>15-17 yrs</td>\n",
       "      <td>224</td>\n",
       "      <td>213574</td>\n",
       "      <td>15.978335</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.373711</td>\n",
       "      <td>0.60</td>\n",
       "      <td>0.09</td>\n",
       "      <td>0.70</td>\n",
       "      <td>1.38</td>\n",
       "      <td>104.881683</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1971</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>all</td>\n",
       "      <td>15-17 yrs</td>\n",
       "      <td>241</td>\n",
       "      <td>220026</td>\n",
       "      <td>15.985843</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.316049</td>\n",
       "      <td>0.66</td>\n",
       "      <td>0.09</td>\n",
       "      <td>0.76</td>\n",
       "      <td>1.52</td>\n",
       "      <td>109.532509</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1972</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>all</td>\n",
       "      <td>15-17 yrs</td>\n",
       "      <td>270</td>\n",
       "      <td>224877</td>\n",
       "      <td>15.985734</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.275120</td>\n",
       "      <td>0.74</td>\n",
       "      <td>0.09</td>\n",
       "      <td>0.78</td>\n",
       "      <td>1.61</td>\n",
       "      <td>120.065636</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1973</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>all</td>\n",
       "      <td>15-17 yrs</td>\n",
       "      <td>258</td>\n",
       "      <td>227256</td>\n",
       "      <td>15.992462</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.200450</td>\n",
       "      <td>0.79</td>\n",
       "      <td>0.10</td>\n",
       "      <td>0.79</td>\n",
       "      <td>1.69</td>\n",
       "      <td>113.528358</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1974</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>all</td>\n",
       "      <td>15-17 yrs</td>\n",
       "      <td>224</td>\n",
       "      <td>229025</td>\n",
       "      <td>16.000776</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.081136</td>\n",
       "      <td>0.83</td>\n",
       "      <td>0.16</td>\n",
       "      <td>0.81</td>\n",
       "      <td>1.80</td>\n",
       "      <td>97.805916</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   year  state  legal1820 dtype      agegr  count     pop        age  legal  \\\n",
       "0  1970      1        0.0   all  15-17 yrs    224  213574  15.978335    0.0   \n",
       "1  1971      1        0.0   all  15-17 yrs    241  220026  15.985843    0.0   \n",
       "2  1972      1        0.0   all  15-17 yrs    270  224877  15.985734    0.0   \n",
       "3  1973      1        0.0   all  15-17 yrs    258  227256  15.992462    0.0   \n",
       "4  1974      1        0.0   all  15-17 yrs    224  229025  16.000776    0.0   \n",
       "\n",
       "   beertaxa  beerpercap  winepercap  spiritpercap  totpercap       mrate  \n",
       "0  1.373711        0.60        0.09          0.70       1.38  104.881683  \n",
       "1  1.316049        0.66        0.09          0.76       1.52  109.532509  \n",
       "2  1.275120        0.74        0.09          0.78       1.61  120.065636  \n",
       "3  1.200450        0.79        0.10          0.79       1.69  113.528358  \n",
       "4  1.081136        0.83        0.16          0.81       1.80   97.805916  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = pd.read_stata(\"../data/deaths.dta\")\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>year</th>\n",
       "      <th>state</th>\n",
       "      <th>legal1820</th>\n",
       "      <th>dtype</th>\n",
       "      <th>agegr</th>\n",
       "      <th>count</th>\n",
       "      <th>pop</th>\n",
       "      <th>age</th>\n",
       "      <th>legal</th>\n",
       "      <th>beertaxa</th>\n",
       "      <th>beerpercap</th>\n",
       "      <th>winepercap</th>\n",
       "      <th>spiritpercap</th>\n",
       "      <th>totpercap</th>\n",
       "      <th>mrate</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1382</th>\n",
       "      <td>1970</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>all</td>\n",
       "      <td>18-20 yrs</td>\n",
       "      <td>292</td>\n",
       "      <td>189770</td>\n",
       "      <td>18.969990</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.373711</td>\n",
       "      <td>0.60</td>\n",
       "      <td>0.09</td>\n",
       "      <td>0.70</td>\n",
       "      <td>1.38</td>\n",
       "      <td>153.870468</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1399</th>\n",
       "      <td>1971</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>all</td>\n",
       "      <td>18-20 yrs</td>\n",
       "      <td>316</td>\n",
       "      <td>195623</td>\n",
       "      <td>18.972181</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.316049</td>\n",
       "      <td>0.66</td>\n",
       "      <td>0.09</td>\n",
       "      <td>0.76</td>\n",
       "      <td>1.52</td>\n",
       "      <td>161.535202</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1416</th>\n",
       "      <td>1972</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>all</td>\n",
       "      <td>18-20 yrs</td>\n",
       "      <td>320</td>\n",
       "      <td>200988</td>\n",
       "      <td>18.974701</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.275120</td>\n",
       "      <td>0.74</td>\n",
       "      <td>0.09</td>\n",
       "      <td>0.78</td>\n",
       "      <td>1.61</td>\n",
       "      <td>159.213486</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1433</th>\n",
       "      <td>1973</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>all</td>\n",
       "      <td>18-20 yrs</td>\n",
       "      <td>291</td>\n",
       "      <td>206550</td>\n",
       "      <td>18.970230</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.200450</td>\n",
       "      <td>0.79</td>\n",
       "      <td>0.10</td>\n",
       "      <td>0.79</td>\n",
       "      <td>1.69</td>\n",
       "      <td>140.885986</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1450</th>\n",
       "      <td>1974</td>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>all</td>\n",
       "      <td>18-20 yrs</td>\n",
       "      <td>306</td>\n",
       "      <td>213839</td>\n",
       "      <td>18.980139</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.081136</td>\n",
       "      <td>0.83</td>\n",
       "      <td>0.16</td>\n",
       "      <td>0.81</td>\n",
       "      <td>1.80</td>\n",
       "      <td>143.098312</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      year  state  legal1820 dtype      agegr  count     pop        age  \\\n",
       "1382  1970      1        0.0   all  18-20 yrs    292  189770  18.969990   \n",
       "1399  1971      1        0.0   all  18-20 yrs    316  195623  18.972181   \n",
       "1416  1972      1        0.0   all  18-20 yrs    320  200988  18.974701   \n",
       "1433  1973      1        0.0   all  18-20 yrs    291  206550  18.970230   \n",
       "1450  1974      1        0.0   all  18-20 yrs    306  213839  18.980139   \n",
       "\n",
       "      legal  beertaxa  beerpercap  winepercap  spiritpercap  totpercap  \\\n",
       "1382    0.0  1.373711        0.60        0.09          0.70       1.38   \n",
       "1399    0.0  1.316049        0.66        0.09          0.76       1.52   \n",
       "1416    0.0  1.275120        0.74        0.09          0.78       1.61   \n",
       "1433    0.0  1.200450        0.79        0.10          0.79       1.69   \n",
       "1450    0.0  1.081136        0.83        0.16          0.81       1.80   \n",
       "\n",
       "           mrate  \n",
       "1382  153.870468  \n",
       "1399  161.535202  \n",
       "1416  159.213486  \n",
       "1433  140.885986  \n",
       "1450  143.098312  "
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = data[\n",
    "    (data['year']<=1983)\n",
    "    & (data['agegr']=='18-20 yrs')\n",
    "    & (data['dtype']=='all')\n",
    "].dropna()\n",
    "\n",
    "data.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Causal Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "According to Equation(5.6) in the book, the reduced form of DID equation (with time trend) is as follows:\n",
    "\n",
    "$Y_{st} = \\alpha + \\delta_{rDID}LEGAL_{st} + \\sum_{k}\\beta_{k}STATE_{ks} + \\sum_{j}\\gamma_{j}YEAR_{jt} + \\sum_{k} \\theta_{k}(STATE_{ks} \\times t) + \\epsilon_{st}$\n",
    "\n",
    "where:\n",
    "* $LEGAL_{st}$ is the treatment variable\n",
    "* $YEAR_{jt}$ is the time-fixed effect\n",
    "* $STATE_{ks}$ is the state-fixed effect\n",
    "* $STATE_{ks} \\times t$ is the time trend\n",
    "\n",
    "In the followings, we are interested in the causal effect to all-death death rate when \"fractional legal\" is controlled, which corresponds to the estimate shown in the first row and third columns of table 5.3, with an estiamte of 10.03 and stand error around 4.92"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `DoWhy` Implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not recommended. The [DID API](https://github.com/py-why/dowhy/issues/635) is still under development."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `statsmodels` Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OLS Regression Results: \n",
      " -- Coef: 10.03 \n",
      " --   SE: 4.92\n"
     ]
    }
   ],
   "source": [
    "formula = 'mrate ~ legal + beertaxa + C(state) * year + C(year)'\n",
    "\n",
    "sm_did_model = smf.ols(\n",
    "    formula=formula,\n",
    "    data=data\n",
    ")\n",
    "\n",
    "# regressing with clustered standard error\n",
    "results = sm_did_model.fit(\n",
    "    cov_type='cluster',\n",
    "    cov_kwds={'groups': data['state']}\n",
    ")\n",
    "\n",
    "print(\n",
    "    \"OLS Regression Results: \\n -- Coef: {:.2f} \\n --   SE: {:.2f}\".format(\n",
    "        results.params['legal'], results.bse['legal']\n",
    "    )\n",
    ")\n",
    "\n",
    "# print(results.summary())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mastering-py-metrics",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
